import torch
from transformers import ViltForQuestionAnswering, ViltImageProcessor, BertTokenizer
from PIL import Image

# 1. 设置路径
model_path = "./local_vilt_model"
image_path = "cat.jpg"  # 替换为你的图片路径

# 2. 加载模型和处理器
model = ViltForQuestionAnswering.from_pretrained(model_path)
image_processor = ViltImageProcessor.from_pretrained(model_path)
tokenizer = BertTokenizer.from_pretrained(model_path)  # 使用本地tokenizer

# 3. 加载图片
image = Image.open(image_path)
# question = "How many cats are there?"
# question = "What is in the picture?"
question = "What is the color of the left cat?"


# 4. 正确的预处理方式（分开处理图像和文本）
# 处理图像
image_inputs = image_processor(image, return_tensors="pt")
# 处理文本
text_inputs = tokenizer(question, return_tensors="pt")

# 5. 合并输入
inputs = {
    "pixel_values": image_inputs.pixel_values,
    "input_ids": text_inputs.input_ids,
    "attention_mask": text_inputs.attention_mask
}

# 6. 模型推理
with torch.no_grad():
    outputs = model(**inputs)
    answer_idx = outputs.logits.argmax(-1).item()

print("Predicted answer:", model.config.id2label[answer_idx])
